[Drain] OnTimer - propagate caused by drain bit up to DoFnRunner by stankiewicz · Pull Request #37012 · apache/beam

stankiewicz · 2025-12-05T17:17:20Z

This pull request enhances the Apache Beam Java SDK's timer mechanism by propagating a causedByDrain flag. This flag allows DoFns to distinguish between timers that fire as part of normal pipeline execution and those that are explicitly triggered during a pipeline draining process. By propagating this information through the core timer interfaces and their implementations, it provides more granular control and context for DoFns reacting to timer events e.g. in future within SDF or ReduceFnRunner.

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

stankiewicz · 2025-12-05T17:18:16Z

For stack delta, you look at relevant commit - ontimer

gemini-code-assist · 2025-12-05T17:18:33Z

Summary of Changes

Hello @stankiewicz, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the Apache Beam Java SDK's timer mechanism by introducing a causedByDrain flag. This flag allows DoFns to distinguish between timers that fire as part of normal pipeline execution and those that are explicitly triggered during a pipeline draining process. By propagating this information through the core timer interfaces and their implementations, it provides more granular control and context for DoFns reacting to timer events, which is crucial for robust pipeline shutdown and data consistency during draining.

Highlights

Propagate 'causedByDrain' Flag: Introduced a boolean causedByDrain parameter to the DoFnRunner.onTimer method and TimerInternals.TimerData to indicate if a timer firing was triggered by a draining operation.
DoFn Context Updates: Added causedByDrain() abstract methods to DoFn.ProcessContext and DoFn.OnTimerContext to expose the drain status to user DoFn logic.
Extensive Integration: Updated numerous DoFnRunner implementations, TimerData factory methods, and Timer utility classes across various runners (Core Java, Flink, Dataflow, Samza, Spark) and SDK components to correctly pass and handle the new causedByDrain flag.
Test Coverage: Modified existing test cases to accommodate the new causedByDrain parameter in TimerData.of and Timer.of calls, ensuring consistency and correctness.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

github-actions · 2025-12-06T00:28:59Z

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

stankiewicz · 2025-12-06T09:18:07Z

Run Java PreCommit

github-actions · 2026-02-04T12:48:18Z

This pull request has been marked as stale due to 60 days of inactivity. It will be closed in 1 week if no further activity occurs. If you think that’s incorrect or this pull request requires a review, please simply write any comment. If closed, you can revive the PR at any time and @mention a reviewer or discuss it on the dev@beam.apache.org list. Thank you for your contributions.

codecov · 2026-02-05T12:58:41Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 40.06%. Comparing base (4e76198) to head (2612296).

Additional details and impacted files

@@            Coverage Diff            @@
##             master   #37012   +/-   ##
=========================================
  Coverage     40.06%   40.06%           
  Complexity     3404     3404           
=========================================
  Files          1177     1177           
  Lines        187083   187083           
  Branches       3581     3581           
=========================================
+ Hits          74947    74953    +6     
+ Misses       108744   108739    -5     
+ Partials       3392     3391    -1

Flag	Coverage Δ
java	`71.93% <ø> (+0.02%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java

runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java

stankiewicz · 2026-02-17T13:13:05Z

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java

                    .setPaneInfo(currentTimer.getPaneInfo())
+                    .setCausedByDrain(causedByDrain)
                    .setReceiver(
                        windowedValue ->


@kennknowles , ptal.

causedBy drain will be part of windowedValue but because receiver is lambda from WV to context.outputWindowedValue, I'm pretty sure I'm loosing all WV metadata.

stankiewicz · 2026-02-17T13:15:39Z

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java

            .setReceiver(
                windowedValue -> {
                  checkTimerTimestamp(windowedValue.getTimestamp());
                  outputTo(mainOutputConsumer, windowedValue);


@kennknowles here, windowedValue is sent differently, why is that?

I am not sure. This file was a huge mess, with actually many types of transforms executed in the same file via switch statements. It may just be an accidental piece of history and/or a part that I failed to refactor when I introduced OutputBuilder. My quick read is that this one is the idea way to do it, because it passes the whole WindowedValue to the FnDataReceiver so it doesn't lose dasta. In the other places where it goes through the OnTimerContext it loses metadata.

stankiewicz · 2026-02-17T13:20:21Z

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java

        Instant timestamp,
        Collection<? extends BoundedWindow> windows,
        PaneInfo paneInfo) {
      builder(value).setTimestamp(timestamp).setWindows(windows).setPaneInfo(paneInfo).output();


@kennknowles , I assume that by introducing builder we freeze outputWindowedValue and have some limited functionality, but as part of receiver we should use outputWindowedValue method as we risk losing any new fields we've added.

I am not sure if I understand the comment. Here is what I think are answers:

Yes, we freeze the API for outputWindowedValue.

No, we shouldn't lose fields, because builder(value) is an abstract method and it is the job of OutputReceiver.builder(...) to make sure to set all the builder values to defaults, and propagate the values from the current element or timer context.

…ions. Mostly passthrough.

github-actions · 2026-02-17T16:25:28Z

Assigning reviewers:

R: @chamikaramj for label java.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

kennknowles · 2026-02-17T17:59:51Z

...n/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java


+    @Override
+    public CausedByDrain causedByDrain() {
+      return CausedByDrain.NORMAL;


Should this come from the element metadata?

kennknowles · 2026-02-17T18:00:23Z

...n/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java


+    @Override
+    public CausedByDrain causedByDrain() {
+      return CausedByDrain.NORMAL;


Should this come from the element metadata?

kennknowles · 2026-02-17T18:01:28Z

.../core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java

+        timestamp,
+        outputTimestamp,
+        timeDomain,
+        CausedByDrain.NORMAL);


I think it should propagate here? will that be a follow-up that adds it to the pushback side input DoFnRunner?

kennknowles · 2026-02-17T18:02:06Z

runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java

        new Instant(0),
-        TimeDomain.EVENT_TIME);
+        TimeDomain.EVENT_TIME,
+        CausedByDrain.NORMAL);


Maybe this unit test file should test caused by drain propagation somehow?

kennknowles · 2026-02-17T18:04:06Z

...ava-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/wire/CommonCoderTest.java

          new Instant(((Number) kvMap.get("holdTimestamp")).longValue()),
-          paneInfo);
+          paneInfo,
+          CausedByDrain.NORMAL);


I think this should actually come from the data? It is a weird method and I have no context for why it is the way it is...

kennknowles · 2026-02-17T18:07:19Z

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java

        Instant timestamp,
        Collection<? extends BoundedWindow> windows,
        PaneInfo paneInfo) {
      builder(value).setTimestamp(timestamp).setWindows(windows).setPaneInfo(paneInfo).output();


I am not sure if I understand the comment. Here is what I think are answers:

Yes, we freeze the API for outputWindowedValue.

No, we shouldn't lose fields, because builder(value) is an abstract method and it is the job of OutputReceiver.builder(...) to make sure to set all the builder values to defaults, and propagate the values from the current element or timer context.

kennknowles · 2026-02-17T18:07:43Z

sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/TimerTest.java

            HOLD_TIME,
-            PaneInfo.NO_FIRING));
+            PaneInfo.NO_FIRING,
+            CausedByDrain.NORMAL));


We need some caused by drain tests?

github-actions · 2026-02-25T12:27:47Z

Reminder, please take a look at this pr: @chamikaramj

github-actions bot added java runners spark dataflow direct core flink fn-execution jet samza labels Dec 5, 2025

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch 2 times, most recently from 0dbde3d to 69b077f Compare December 5, 2025 21:33

github-actions bot added the stale label Feb 4, 2026

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch from 69b077f to 9ff470a Compare February 5, 2026 09:14

github-actions bot removed the direct label Feb 5, 2026

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch 2 times, most recently from 2d77410 to 2612296 Compare February 5, 2026 11:04

github-actions bot removed the stale label Feb 5, 2026

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch from 2612296 to bf072d2 Compare February 5, 2026 13:30

stankiewicz commented Feb 5, 2026

View reviewed changes

runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java Outdated Show resolved Hide resolved

stankiewicz commented Feb 5, 2026

View reviewed changes

runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java Outdated Show resolved Hide resolved

stankiewicz commented Feb 5, 2026

View reviewed changes

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java Show resolved Hide resolved

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch from bf072d2 to 602243b Compare February 17, 2026 12:54

stankiewicz commented Feb 17, 2026

View reviewed changes

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch 2 times, most recently from c68df92 to 43050a8 Compare February 17, 2026 13:59

add causedByDrain to DoFnRunner.onTimer interface and all implementat…

03aa66f

…ions. Mostly passthrough.

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch from 43050a8 to 03aa66f Compare February 17, 2026 14:23

github-actions bot added the Next Action: Reviewers label Feb 17, 2026

stankiewicz requested a review from kennknowles February 17, 2026 16:37

kennknowles requested changes Feb 17, 2026

View reviewed changes

github-actions bot added the slow-review label Feb 25, 2026

Comments

Conversation

stankiewicz commented Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

GitHub Actions Tests Status (on master branch)

Uh oh!

stankiewicz commented Dec 5, 2025

Uh oh!

gemini-code-assist bot commented Dec 5, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

github-actions bot commented Dec 6, 2025

Uh oh!

stankiewicz commented Dec 6, 2025

Uh oh!

github-actions bot commented Feb 4, 2026

Uh oh!

codecov bot commented Feb 5, 2026

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Feb 17, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Feb 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

stankiewicz commented Dec 5, 2025 •

edited

Loading